# -*- coding: utf-8 -*-
"""
超跌选股筛选器
筛选规则：
1. CCI 指标进入超卖区，即 CCI数值跌破-100时
2. KDJ数值小于30
3. WR 指标突破 80的水平线
4. RSI 指标小于30
5. 布林带：收盘价跌破布林下轨
6. 乖离率：BIAS(10日) < -10%

超跌信号说明：
- CCI < -100: 商品通道指数进入超卖区，价格可能被低估
- KDJ < 30: 随机指标显示超卖，可能即将反弹
- WR > 80: 威廉指标显示超卖（WR越大越超卖）
- RSI < 30: 相对强弱指数显示超卖，可能反弹
- 收盘价 < 布林下轨: 价格严重偏离正常区间，超卖
- BIAS < -10%: 价格严重偏离均线，超卖

使用方法：
    python oversold_screener.py              # 筛选所有股票
    python oversold_screener.py --limit 100  # 限制筛选数量（测试用）
    python oversold_screener.py --stock 000001  # 筛选指定股票
"""
import logging
import pandas as pd
import numpy as np
from datetime import datetime
import sys
import os
import argparse

# 添加父目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import DatabaseManager, StockInfoDAO
from config import DATABASE_CONFIG
import warnings

# 忽略pandas警告
warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)

# 配置日志
# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'oversold_screener.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


class OversoldScreener:
    """超跌选股筛选器"""
    
    def __init__(self):
        self.db = DatabaseManager(DATABASE_CONFIG)
        self.stock_dao = StockInfoDAO(self.db)
    
    def get_stock_codes(self, limit: int = None, specific_stock: str = None) -> list:
        """获取所有活跃股票代码"""
        try:
            if specific_stock:
                sql = "SELECT stock_code, stock_name FROM new_stock_info WHERE stock_code = %s AND status = 'active'"
                results = self.db.execute_query(sql, (specific_stock,))
            else:
                sql = "SELECT stock_code, stock_name FROM new_stock_info WHERE status = 'active' ORDER BY stock_code"
                if limit:
                    sql += f" LIMIT {limit}"
                results = self.db.execute_query(sql)
            
            return [(row['stock_code'], row['stock_name']) for row in results]
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return []
    
    def get_stock_daily_data(self, stock_code: str, days: int = 50) -> pd.DataFrame:
        """
        获取股票最近N天的日线数据
        
        Args:
            stock_code: 股票代码
            days: 获取天数（默认50天，确保有足够数据计算指标）
            
        Returns:
            包含交易数据的DataFrame
        """
        sql = """
        SELECT trade_date, open_price, high_price, low_price, close_price, 
               volume, amount, change_pct, turnover_rate
        FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT %s
        """
        try:
            results = self.db.execute_query(sql, (stock_code, days))
            if not results:
                return pd.DataFrame()
            
            df = pd.DataFrame(results)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 转换数值列
            for col in ['open_price', 'high_price', 'low_price', 'close_price', 
                       'volume', 'amount', 'change_pct', 'turnover_rate']:
                df[col] = pd.to_numeric(df[col], errors='coerce')
            
            # 按日期升序排列
            df = df.sort_values('trade_date').reset_index(drop=True)
            
            return df
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 日线数据失败: {e}")
            return pd.DataFrame()
    
    @staticmethod
    def calculate_cci(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """
        计算CCI指标（Commodity Channel Index - 商品通道指数）
        
        CCI = (TP - MA) / (0.015 * MD)
        其中：
        - TP (Typical Price) = (最高价 + 最低价 + 收盘价) / 3
        - MA = TP的N日简单移动平均
        - MD = TP与MA的平均绝对偏差
        
        Args:
            df: 包含high_price, low_price, close_price的DataFrame
            period: 计算周期，默认14
            
        Returns:
            添加了cci列的DataFrame
        """
        if df.empty or len(df) < period:
            df['cci'] = np.nan
            return df
        
        # 计算典型价格
        tp = (df['high_price'] + df['low_price'] + df['close_price']) / 3
        
        # 计算TP的移动平均
        ma = tp.rolling(window=period).mean()
        
        # 计算平均绝对偏差
        md = tp.rolling(window=period).apply(lambda x: np.abs(x - x.mean()).mean())
        
        # 计算CCI
        df['cci'] = (tp - ma) / (0.015 * md)
        
        return df
    
    @staticmethod
    def calculate_wr(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """
        计算WR指标（Williams %R - 威廉指标）
        
        WR = (Hn - C) / (Hn - Ln) * 100
        其中：
        - Hn = N日内最高价
        - Ln = N日内最低价
        - C = 当日收盘价
        
        注意：WR的值域在0到100之间，值越大表示越超卖
        - WR > 80: 超卖区（可能反弹）
        - WR < 20: 超买区（可能回调）
        
        Args:
            df: 包含high_price, low_price, close_price的DataFrame
            period: 计算周期，默认14
            
        Returns:
            添加了wr列的DataFrame
        """
        if df.empty or len(df) < period:
            df['wr'] = np.nan
            return df
        
        # 计算N日内最高价和最低价
        highest_high = df['high_price'].rolling(window=period).max()
        lowest_low = df['low_price'].rolling(window=period).min()
        
        # 计算WR
        df['wr'] = (highest_high - df['close_price']) / (highest_high - lowest_low + 0.0001) * 100
        
        return df
    
    @staticmethod
    def calculate_rsi(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """
        计算RSI相对强弱指数（Relative Strength Index）
        
        RSI = 100 - (100 / (1 + RS))
        其中 RS = 平均上涨幅度 / 平均下跌幅度
        
        Args:
            df: 包含close_price列的DataFrame
            period: 周期，默认14
            
        Returns:
            添加了rsi列的DataFrame
        """
        if df.empty or len(df) < period:
            df['rsi'] = np.nan
            return df
        
        # 计算价格变化
        delta = df['close_price'].diff()
        
        # 分离上涨和下跌
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        
        # 计算平均涨幅和跌幅
        avg_gain = gain.rolling(window=period).mean()
        avg_loss = loss.rolling(window=period).mean()
        
        # 避免除零
        avg_loss = avg_loss.replace(0, 0.0001)
        
        # 计算RS和RSI
        rs = avg_gain / avg_loss
        df['rsi'] = 100 - (100 / (1 + rs))
        
        return df
    
    @staticmethod
    def calculate_bollinger(df: pd.DataFrame, period: int = 20, std_dev: float = 2.0) -> pd.DataFrame:
        """
        计算布林带（Bollinger Bands）
        
        中轨 = MA(20)
        上轨 = 中轨 + 2 * 标准差
        下轨 = 中轨 - 2 * 标准差
        
        Args:
            df: 包含close_price列的DataFrame
            period: 周期，默认20
            std_dev: 标准差倍数，默认2.0
            
        Returns:
            添加了boll_mid, boll_upper, boll_lower列的DataFrame
        """
        if df.empty or len(df) < period:
            df['boll_mid'] = np.nan
            df['boll_upper'] = np.nan
            df['boll_lower'] = np.nan
            return df
        
        # 中轨（移动平均）
        df['boll_mid'] = df['close_price'].rolling(window=period).mean()
        
        # 标准差
        std = df['close_price'].rolling(window=period).std()
        
        # 上轨和下轨
        df['boll_upper'] = df['boll_mid'] + std_dev * std
        df['boll_lower'] = df['boll_mid'] - std_dev * std
        
        return df
    
    @staticmethod
    def calculate_bias(df: pd.DataFrame, period: int = 10) -> pd.DataFrame:
        """
        计算乖离率（BIAS）
        
        BIAS = (当前价 - N日均价) / N日均价 * 100%
        
        乖离率衡量股价偏离移动平均线的程度
        
        Args:
            df: 包含close_price列的DataFrame
            period: 周期，默认10日
            
        Returns:
            添加了bias列的DataFrame
        """
        if df.empty or len(df) < period:
            df['bias'] = np.nan
            return df
        
        # 计算移动平均
        ma = df['close_price'].rolling(window=period).mean()
        
        # 计算乖离率
        df['bias'] = (df['close_price'] - ma) / ma * 100
        
        return df
    
    @staticmethod
    def calculate_kdj(df: pd.DataFrame, n: int = 9, m1: int = 3, m2: int = 3) -> pd.DataFrame:
        """
        计算KDJ指标（随机指标）
        
        Args:
            df: 包含high_price, low_price, close_price的DataFrame
            n: RSV周期，默认9
            m1: K值周期，默认3
            m2: D值周期，默认3
            
        Returns:
            添加了kdj_k, kdj_d, kdj_j列的DataFrame
        """
        if df.empty or len(df) < n:
            df['kdj_k'] = np.nan
            df['kdj_d'] = np.nan
            df['kdj_j'] = np.nan
            return df
        
        # 计算RSV
        low_min = df['low_price'].rolling(window=n).min()
        high_max = df['high_price'].rolling(window=n).max()
        
        # 避免除零
        rsv = 100 * (df['close_price'] - low_min) / (high_max - low_min + 0.0001)
        
        # 计算K值
        df['kdj_k'] = rsv.ewm(span=m1, adjust=False).mean()
        
        # 计算D值
        df['kdj_d'] = df['kdj_k'].ewm(span=m2, adjust=False).mean()
        
        # 计算J值
        df['kdj_j'] = 3 * df['kdj_k'] - 2 * df['kdj_d']
        
        return df
    
    def check_oversold_conditions(self, df: pd.DataFrame) -> dict:
        """
        检查是否满足超跌条件
        
        条件：
        1. CCI < -100 (进入超卖区)
        2. KDJ < 30 (超卖)
        3. WR > 80 (超卖)
        4. RSI < 30 (超卖)
        5. 收盘价 < 布林下轨 (超卖)
        6. BIAS < -10% (超卖)
        
        Args:
            df: 包含技术指标的DataFrame
            
        Returns:
            包含检查结果的字典
        """
        if df.empty or len(df) < 2:
            return {
                'is_oversold': False,
                'reason': '数据不足',
                'latest_date': None,
                'cci': None,
                'kdj': None,
                'wr': None,
                'rsi': None,
                'boll_lower': None,
                'bias': None,
                'close_price': None
            }
        
        # 获取最新一天的数据
        latest = df.iloc[-1]
        
        # 检查是否有有效数据
        required_fields = ['cci', 'kdj_k', 'wr', 'rsi', 'boll_lower', 'bias', 'close_price']
        if any(pd.isna(latest[field]) for field in required_fields):
            return {
                'is_oversold': False,
                'reason': '指标数据缺失',
                'latest_date': latest['trade_date'],
                'cci': latest['cci'],
                'kdj': latest['kdj_k'],
                'wr': latest['wr'],
                'rsi': latest['rsi'],
                'boll_lower': latest['boll_lower'],
                'bias': latest['bias'],
                'close_price': latest['close_price']
            }
        
        # 检查超跌条件
        cci_oversold = latest['cci'] < -100
        kdj_oversold = latest['kdj_k'] < 30
        wr_oversold = latest['wr'] > 80
        rsi_oversold = latest['rsi'] < 30
        boll_oversold = latest['close_price'] < latest['boll_lower']
        bias_oversold = latest['bias'] < -10
        
        is_oversold = (cci_oversold and kdj_oversold and wr_oversold and 
                      rsi_oversold and boll_oversold and bias_oversold)
        
        # 构建原因说明
        conditions_met = []
        if cci_oversold:
            conditions_met.append(f"CCI={latest['cci']:.2f}<-100")
        if kdj_oversold:
            conditions_met.append(f"KDJ={latest['kdj_k']:.2f}<30")
        if wr_oversold:
            conditions_met.append(f"WR={latest['wr']:.2f}>80")
        if rsi_oversold:
            conditions_met.append(f"RSI={latest['rsi']:.2f}<30")
        if boll_oversold:
            conditions_met.append(f"价格={latest['close_price']:.2f}<布林下轨={latest['boll_lower']:.2f}")
        if bias_oversold:
            conditions_met.append(f"BIAS={latest['bias']:.2f}%<-10%")
        
        reason = ', '.join(conditions_met) if conditions_met else '不满足超跌条件'
        
        return {
            'is_oversold': is_oversold,
            'reason': reason,
            'latest_date': latest['trade_date'],
            'cci': latest['cci'],
            'kdj': latest['kdj_k'],
            'wr': latest['wr'],
            'rsi': latest['rsi'],
            'boll_lower': latest['boll_lower'],
            'bias': latest['bias'],
            'close_price': latest['close_price'],
            'cci_met': cci_oversold,
            'kdj_met': kdj_oversold,
            'wr_met': wr_oversold,
            'rsi_met': rsi_oversold,
            'boll_met': boll_oversold,
            'bias_met': bias_oversold
        }
    
    def screen_single_stock(self, stock_code: str, stock_name: str) -> dict:
        """
        筛选单只股票
        
        Args:
            stock_code: 股票代码
            stock_name: 股票名称
            
        Returns:
            筛选结果字典
        """
        try:
            # 获取日线数据
            df = self.get_stock_daily_data(stock_code, days=50)
            
            if df.empty:
                return {
                    'stock_code': stock_code,
                    'stock_name': stock_name,
                    'is_oversold': False,
                    'reason': '无日线数据'
                }
            
            # 计算技术指标
            df = self.calculate_cci(df, period=14)
            df = self.calculate_kdj(df, n=9, m1=3, m2=3)
            df = self.calculate_wr(df, period=14)
            df = self.calculate_rsi(df, period=14)
            df = self.calculate_bollinger(df, period=20, std_dev=2.0)
            df = self.calculate_bias(df, period=10)
            
            # 检查超跌条件
            result = self.check_oversold_conditions(df)
            result['stock_code'] = stock_code
            result['stock_name'] = stock_name
            
            return result
            
        except Exception as e:
            logger.error(f"筛选股票 {stock_code} 失败: {e}")
            return {
                'stock_code': stock_code,
                'stock_name': stock_name,
                'is_oversold': False,
                'reason': f'筛选失败: {e}'
            }
    
    def screen_stocks(self, limit: int = None, specific_stock: str = None) -> list:
        """
        批量筛选股票
        
        Args:
            limit: 限制筛选数量
            specific_stock: 指定股票代码
            
        Returns:
            筛选结果列表
        """
        # 获取股票列表
        stock_list = self.get_stock_codes(limit=limit, specific_stock=specific_stock)
        
        if not stock_list:
            logger.warning("没有找到符合条件的股票")
            return []
        
        logger.info(f"开始筛选 {len(stock_list)} 只股票...")
        
        results = []
        oversold_stocks = []
        
        for i, (stock_code, stock_name) in enumerate(stock_list, 1):
            # 显示进度
            if i % 100 == 0 or i == len(stock_list):
                logger.info(f"进度: {i}/{len(stock_list)}")
            
            # 筛选单只股票
            result = self.screen_single_stock(stock_code, stock_name)
            
            if result['is_oversold']:
                oversold_stocks.append(result)
                logger.info(f"✓ 发现超跌股票: {stock_code} {stock_name} - {result['reason']}")
            
            results.append(result)
        
        logger.info(f"筛选完成！共发现 {len(oversold_stocks)} 只超跌股票")
        
        return results
    
    def save_results(self, results: list, output_format: str = 'both'):
        """
        保存筛选结果
        
        Args:
            results: 筛选结果列表
            output_format: 输出格式 'csv', 'txt', 'both'
        """
        # 过滤出超跌股票
        oversold_stocks = [r for r in results if r['is_oversold']]
        
        if not oversold_stocks:
            logger.info("没有发现超跌股票，不生成输出文件")
            return
        
        # 生成时间戳
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        
        # 输出逗号分隔的股票代码
        stock_codes = [s['stock_code'] for s in oversold_stocks]
        codes_string = ','.join(stock_codes)
        
        print("\n" + "=" * 80)
        print("超跌股票代码（逗号分隔）：")
        print("=" * 80)
        print(codes_string)
        print("=" * 80)
        
        # 保存CSV文件
        if output_format in ['csv', 'both']:
            df = pd.DataFrame(oversold_stocks)
            csv_file = f'oversold_stocks_{timestamp}.csv'
            df.to_csv(csv_file, index=False, encoding='utf-8-sig')
            logger.info(f"CSV文件已保存: {csv_file}")
        
        # 保存TXT文件（股票代码列表）
        if output_format in ['txt', 'both']:
            txt_file = f'stock_codes_oversold_{timestamp}.txt'
            with open(txt_file, 'w', encoding='utf-8') as f:
                f.write('\n'.join(stock_codes))
            logger.info(f"股票代码文件已保存: {txt_file}")
        
        # 打印详细信息
        print("\n超跌股票详细信息：")
        print("-" * 170)
        print(f"{'股票代码':<10} {'股票名称':<10} {'最新日期':<12} {'CCI':<8} {'KDJ':<8} {'WR':<8} {'RSI':<8} {'BIAS%':<8} {'收盘价':<10} {'布林下轨':<10}")
        print("-" * 170)
        
        for stock in oversold_stocks:
            print(f"{stock['stock_code']:<10} {stock['stock_name']:<10} "
                  f"{str(stock['latest_date'])[:10]:<12} "
                  f"{stock['cci']:<8.2f} {stock['kdj']:<8.2f} {stock['wr']:<8.2f} "
                  f"{stock['rsi']:<8.2f} {stock['bias']:<8.2f} "
                  f"{stock['close_price']:<10.2f} {stock['boll_lower']:<10.2f}")
        
        print("-" * 170)


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='超跌选股筛选器')
    parser.add_argument('--limit', type=int, help='限制筛选股票数量（用于测试）')
    parser.add_argument('--stock', type=str, help='指定筛选某只股票代码')
    parser.add_argument('--format', type=str, default='both', 
                       choices=['csv', 'txt', 'both'],
                       help='输出格式：csv, txt, both (默认both)')
    
    args = parser.parse_args()
    
    try:
        logger.info("=" * 80)
        logger.info("超跌选股筛选器启动")
        logger.info("筛选条件: CCI < -100 AND KDJ < 30 AND WR > 80 AND RSI < 30")
        logger.info("         AND 收盘价 < 布林下轨 AND BIAS < -10%")
        logger.info("=" * 80)
        
        # 创建筛选器
        screener = OversoldScreener()
        
        # 执行筛选
        results = screener.screen_stocks(limit=args.limit, specific_stock=args.stock)
        
        # 保存结果
        screener.save_results(results, output_format=args.format)
        
        logger.info("=" * 80)
        logger.info("筛选完成")
        logger.info("=" * 80)
        
    except Exception as e:
        logger.error(f"程序执行失败: {e}", exc_info=True)
        return 1
    
    return 0


if __name__ == "__main__":
    sys.exit(main())

